Proto: add DataSink serialization hook - #23752
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #23752 +/- ##
========================================
Coverage 80.85% 80.85%
========================================
Files 1099 1100 +1
Lines 374304 374416 +112
Branches 374304 374416 +112
========================================
+ Hits 302642 302743 +101
- Misses 53607 53613 +6
- Partials 18055 18060 +5 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
@kumarUjjawal would you be able to work with @Phoenix500526 to review this and coordinate with your work to avoid conflicts? |
| Ok(LexRequirement::new(sort_exprs.into_iter().map(Into::into))) | ||
| } | ||
|
|
||
| fn partitioned_file_to_proto( |
There was a problem hiding this comment.
#23683 now provides these shared PartitionedFile helpers. Once that gets merged in the main we can merge main to this pr so we can remove this duplicate code, and reuse those helpers.
There was a problem hiding this comment.
OK, I'll rebase this PR when #23683 is merged.
3c3ad39 to
91da434
Compare
…on-datasource (apache#24006) ## Which issue does this PR close? - Part of apache#23494. Precursor for apache#23497 / apache#23683 (`DataSource` / `FileSource` proto hooks) and for apache#23752. ## Rationale for this change The protobuf conversions for the file-scan leaf types — `PartitionedFile`, `FileGroup`, `FileRange` — live in `datafusion-proto` as `TryFromProto` impls, because that is historically the only crate that can name both sides (the DataFusion type and the prost message are both foreign to it, hence the `TryFromProto` workaround trait in the first place). That placement means any *other* crate that needs those conversions has to reimplement them. apache#23683 hits exactly this: a `FileSource` serializing its own scan config needs to encode file groups, so the first cut of that PR grew a private second copy of the `PartitionedFile` wire logic inside `datafusion-datasource`, which can then drift from the central serializer. The same will be true of every source migrated under apache#23516–apache#23518. Nothing about these conversions needs `datafusion-proto`: they are plain data, with `ScalarValue` / `Statistics` / `Schema` going through `datafusion-proto-common`. They belong next to the types. ## What changes are included in this PR? - New `datafusion_datasource::proto` module, behind a new `proto` feature on `datafusion-datasource` (off by default; `datafusion-proto` enables it): - `FileRange::try_to_proto` / `try_from_proto` - `PartitionedFile::try_to_proto` / `try_from_proto` - `FileGroup` <-> `protobuf::FileGroup` - `datafusion-proto`'s `TryFromProto` impls for those types become one-line shims delegating to the new impls, so every existing caller keeps working and the two sides cannot disagree. ### Why these are `TryFrom` and not `try_to_proto` hooks `TryFromProto` exists because `datafusion-proto` owns neither side of the conversions it hosts: with both the DataFusion type and the prost message foreign to it, `impl TryFrom<protobuf::X> for X` is rejected by the orphan rule, so a local trait was the only way to say the same thing. Moving a conversion into the crate that owns the DataFusion type removes that constraint, and `&T` is `#[fundamental]`, so both directions are expressible with the standard trait (checked, not assumed): ```rust impl TryFrom<&protobuf::PartitionedFile> for PartitionedFile // ok impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile // ok impl TryFrom<&[PartitionedFile]> for protobuf::FileGroup // E0117 ``` The last one is why `protobuf::FileGroup`'s *slice* conversion stays a `TryFromProto` shim: `&[PartitionedFile]` is not a type this crate owns, while `&FileGroup` is. Callers inside DataFusion go through `FileGroup`. So the rule this PR sets for the rest of apache#23494: **plain data uses `TryFrom`; anything needing an encode/decode context keeps the `try_to_proto(ctx)` / `try_from_proto(node, ctx)` hooks**, because the standard trait cannot carry that second argument. Usefully, none of the ~40 `TryFromProto`/`FromProto` impls needs a context, and nothing that needs one was ever a `TryFromProto` impl — the two categories are already disjoint, so the shape now tells a reader whether a conversion recurses. ### Why now `FromProto` / `TryFromProto` were added in apache#21929, *after* the 54.0.0 release, and 54.1.0 was cut before any of this landed — so they have never shipped in a release. Replacing them with the standard traits, and eventually deleting them, is a no-op for semver **today** and a major breaking change the moment 55.0.0 goes out. The same applies to the six inherent `try_to_proto` / `try_from_proto` methods this PR would otherwise have added: they are new, unreleased API, so choosing their final shape costs nothing right now. The other reason to settle it here rather than in a follow-up: this is the PR that establishes the pattern for the data-source family (apache#23516-apache#23519 and apache#23752 / apache#23781 are all queued behind it). Whichever shape merges first is the one they will copy. Retiring the remaining ~34 impls is still its own follow-up. Two notes for whoever picks it up: the sink and format-option conversions can move next to their types the same way, but the ones for `datafusion-common`-owned types (`JoinType`, `NullEquality`, `TableReference`, `UnnestOptions`, ...) cannot — `datafusion-common` cannot depend on `datafusion-proto-models` (it is underneath it via `datafusion-proto-common`). Their legal home is `proto-models` itself, implementing on the local proto type, which is already how `proto-common` hosts the `ScalarValue` / `Statistics` conversions. ## Are these changes tested? Yes. - New unit tests in `datafusion_datasource::proto` covering the `PartitionedFile` round trip (path, size, mtime, partition values, range, arrow schema, statistics), the `FileGroup` round trip, and the invalid-path error. - The existing `datafusion-proto` tests now exercise the delegating shims, so they also pin the shims themselves. - `datafusion-proto`, all features: 227 passed / 0 failed. - `datafusion-datasource` with `proto`: 180 passed / 0 failed. - Full workspace run: 10347 passed (`cargo test --profile ci --workspace --lib --tests --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption`). The only failures are 8 backtrace-symbolization tests in `datafusion-common`, a crate this PR does not touch and which sits below every crate it does; they fail the same way on the base commit on macOS. - `cargo fmt` and `ci/scripts/rust_clippy.sh` clean. ## Are there any user-facing changes? The protobuf wire format is unchanged, and no existing API changes shape. Additive: - New `proto` feature on `datafusion-datasource` (off by default). - New `TryFrom` impls in both directions between `FileRange`, `PartitionedFile`, `FileGroup` and their protobuf messages, under that feature. No new names are added to the crate's API surface: the trait is `core::convert::TryFrom`. Note for reviewers: while writing the round-trip test I found that `PartitionedFile` statistics do not round-trip cleanly on `main` — filed as apache#23998. This PR preserves that behavior exactly rather than changing decode semantics in a refactor; the test documents it. --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
|
Thanks for this @Phoenix500526 — the Rebase first. #24006 merged and touches the same two places this PR does, so it now conflicts in
Please use impl TryFrom<&FileSinkConfig> for protobuf::FileSinkConfig { ... }
impl TryFrom<&protobuf::FileSinkConfig> for FileSinkConfig { ... }with
The hook signature is worth reconsidering before three sinks depend on it. fn try_to_proto(&self, input: PhysicalPlanNode, sort_order: Option<PhysicalSortExprNodeCollection>, sink_schema: &Schema)Compare
Passing the ctx instead fixes both: fn try_to_proto(&self, exec: &DataSinkExec, ctx: &ExecutionPlanEncodeCtx<'_>) -> Result<Option<PhysicalPlanNode>>with a shared helper for the sort-order encoding so the three sinks do not repeat it. If you would rather not restructure, merging this together with #23781 also closes the double-encode window.
Make the encode side of One test no longer asserts anything. In let encoded = config.to_proto()?;
let compatibility_encoded = protobuf::FileSinkConfig::try_from_proto(&config)?;
assert_eq!(encoded, compatibility_encoded);Since Behavior change worth stating in the description. Decode now errors on unknown Finally, could you resolve the review threads you have already addressed? Two are marked done but still open, which makes it hard to see what is left. |
DataSinkExec needs a generic way to delegate protobuf encoding to its wrapped sink. The new hook keeps sink-specific wire logic out of the central serializer while retaining the fallback for unmigrated sinks. Add focused coverage for child and sort-order encoding and verify that existing file sinks continue to use the compatibility path. Refs apache#23498
FileSinkConfig is owned by datafusion-datasource, but its protobuf conversion lived in the central proto crate. Co-locating it lets sink implementations reuse the wire logic without depending on datafusion-proto. Keep existing TryFromProto implementations as compatibility delegates and verify both APIs produce the same protobuf data. CLOSES apache#23498
Protobuf enum accessors silently replace unknown values with defaults. Rejecting malformed values prevents serialized plans from changing sink behavior and verifies compatibility decoding against the encoded config. Refs apache#23498
Defer encoding until a sink opts into serialization so the legacy fallback does not encode child plans and sort expressions twice. Restore FileSinkConfig's standard TryFrom API and map enum values by name to preserve wire compatibility independently of discriminants. Refs apache#23498
91da434 to
99a7f62
Compare
Done |
|
Thanks and sorry about the AI comment - I didn't mean for it to get posted. |
apache#24003) ## Which issue does this PR close? - Part of apache#23494 (the `try_to_proto` / `try_from_proto` migration). Precursor for apache#23497 / apache#23683 / apache#23498 and for the remaining per-plan migrations that need to encode an output partitioning or an ordering. ## Rationale for this change `Partitioning`'s protobuf conversion currently exists in three places: - inline in `RepartitionExec::try_to_proto`, - inline in `RepartitionExec::try_from_proto`, - in `datafusion-proto`'s `serialize_partitioning` / `parse_protobuf_partitioning`. Every plan or data source migrated to the hooks that has an output partitioning has so far copied it again — apache#23683 is about to add a fourth copy for `FileScanConfig`. The flat `PhysicalSortExprNode` encoding is the same story, one level down and more widespread: `AggregateExec`'s ordering requirement, `SymmetricHashJoinExec`'s left and right sort expressions, the window expressions, and range partitioning each hand-roll the same `map` / `collect` over `PhysicalSortExprNode { expr, asc, nulls_first }`, on both the encode and the decode side. apache#23683 (`output_ordering`) and apache#23752 / apache#23781 (the sinks' required ordering) are about to add two more. There is no reason for this logic to live in the callers: it converts a `Partitioning` (and the `PhysicalSortExpr`s and `ScalarValue`s inside it), and needs nothing from the plan level beyond the ability to encode a child expression. ## What changes are included in this PR? Put the single copy next to the types that own it, taking the expression-level context (`datafusion-physical-expr` and `-common` already carry the `proto` feature): - `PhysicalSortExpr::try_to_proto` / `try_from_proto` (`physical-expr-common`) - `sort_exprs_try_to_proto` / `sort_exprs_try_from_proto` (`physical-expr-common`), the sequence form every caller actually needs - `Partitioning::try_to_proto` / `try_from_proto` (`physical-expr`) So plan hooks can reach them, `ExecutionPlanEncodeCtx` / `ExecutionPlanDecodeCtx` now back the expression-level contexts (`PhysicalExprEncode` / `PhysicalExprDecode`) and hand one out via `expr_ctx()`. That bridge is useful beyond partitioning: from here on any plan hook can pass its ctx straight to an expression-level conversion, which is the shape the rest of the migration wants. The sequence encoder is generic over `Borrow<PhysicalSortExpr>`, so one function serves a `LexOrdering`, a `&[PhysicalSortExpr]`, and a `LexRequirement` mapped through `PhysicalSortExpr::from`. The decoder returns the expressions rather than a `LexOrdering`, because callers disagree on what an empty list means: "no ordering declared" for a scan, an error for an operator that requires one. Routed through the new methods: - `RepartitionExec` and `datafusion-proto`'s central serializer, which also retires `serialize_range_partitioning`, `serialize_range_split_point`, `parse_protobuf_range_partitioning` and `parse_protobuf_range_split_point`, - `AggregateExec`'s ordering requirement, `SymmetricHashJoinExec`'s left/right sort expressions, and the window expressions — encode and decode each. Net: −380 / +458 lines with the new tests included; production code shrinks. The next operator that needs partitioning or ordering serde writes one line instead of sixty. ## Are these changes tested? Yes. - New unit tests for the sequence helpers (option and order fidelity, owned `LexRequirement` input, encode-error propagation, missing inner expression), using the existing `proto_test_util` stubs. - The existing round-trip suites cover the rest, and now exercise the shared path: all four partitioning variants (`datafusion-proto`'s `roundtrip_physical_plan` tests exercise the central serializer, `RepartitionExec`'s own hook tests exercise the plan path), plus aggregate `ORDER BY`, window `ORDER BY` and symmetric-hash-join sort expressions for the ordering helpers. `datafusion-proto`: 209 passed / 0 failed. - Lib suites for the three changed crates: 1613 + 1651 + 80 passed / 0 failed. - Full workspace run: 10344 passed (`cargo test --profile ci --workspace --lib --tests --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption`). The only failures are 8 backtrace-symbolization tests in `datafusion-common` and 4 `datafusion-cli` tests that hard-code a repo-relative `parquet-testing/` path; both are artifacts of running from a linked worktree on macOS, in crates this PR does not touch, and the `datafusion-cli` four pass once that path resolves. - `cargo fmt` and `ci/scripts/rust_clippy.sh` clean. ## Are there any user-facing changes? The protobuf wire format is unchanged. Additive API: - `Partitioning::try_to_proto` / `try_from_proto`, `PhysicalSortExpr::try_to_proto` / `try_from_proto`, `sort_exprs_try_to_proto` / `sort_exprs_try_from_proto` (feature `proto`). - `ExecutionPlanEncodeCtx::expr_ctx()` / `ExecutionPlanDecodeCtx::expr_ctx(schema)`. Two behavior differences, both in error paths: - Out-of-range partition counts now return an error instead of wrapping (`as usize`) or panicking (`try_into().unwrap()` in `parse_protobuf_hash_partitioning`). - A missing sort-expression child now reports which field is missing (`PhysicalSortExpr is missing required field 'expr'`) instead of `Unexpected empty physical expression`, and the same message now replaces the three bespoke ones in `AggregateExec`, `SymmetricHashJoinExec` and the window expressions. Four private helpers in `datafusion-proto` are removed (`serialize_range_partitioning`, `serialize_range_split_point`, `parse_protobuf_range_partitioning`, `parse_protobuf_range_split_point`); the public `serialize_partitioning` / `parse_protobuf_partitioning` keep their signatures and behavior. --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Which issue does this PR close?
Rationale for this change
DataSinkExecserialization currently depends on central downcasting indatafusion-proto. This prevents individualDataSinkimplementations fromowning their protobuf serialization logic.
Additionally,
FileSinkConfigis owned bydatafusion-datasource, while itsprotobuf conversion was implemented in
datafusion-proto. Moving theconversion alongside the type allows file sink implementations to reuse it
without depending on the central proto crate.
This provides the foundation for migrating CSV, JSON, and Parquet sink
serialization in follow-up work.
What changes are included in this PR?
DataSink::try_to_protohook with a defaultimplementation returning
None.DataSinkExec::try_to_protoencode its input and required orderingbefore delegating to the underlying sink.
sinks that have not implemented the hook.
FileSinkConfigprotobuf encoding and decoding intodatafusion-datasource.TryFromProtoimplementations as compatibilitydelegates.
Are these changes tested?
Yes.
DataSinktest verifying thatDataSinkExecdelegatesserialization with the encoded input and required ordering.
FileSinkConfigconversion andcompatibility APIs produce the same protobuf data.
cargo fmt --all.datafusion-protointegration tests.Are there any user-facing changes?
No. This is an internal serialization refactor. Existing protobuf data and
the compatibility serialization path remain supported.